/* benchmark guru interfaces */
#include <Configuration.h>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <Message/OpenMP.h>
#include <Utilities/Timer.h>
#include <benchmark/fft_bench.h>
#include <bitset>
using namespace APPNAMESPACE;

/** running FFT tests */
template<typename T, unsigned FFTENG> 
struct test_all
{

  typedef  T real_type;
  typedef  std::complex<T> complex_type;
  int niters;
  std::string lib_name;

  //constructor
  test_all(const string& aname, int steps=10):lib_name(aname),niters(steps){}

  /** add names 
   *
   * -R1 real-single
   * -RM real-multiple
   * -C1 complex-single
   * -CM complex-multiple
   */
  void add2table(std::vector<std::string>& fft_names)
  {
    fft_names.push_back(lib_name+"-R1");
    fft_names.push_back(lib_name+"-RM");
    fft_names.push_back(lib_name+"-C1");
    fft_names.push_back(lib_name+"-CM");
  }

  /* perform the FFT with the given engine 
   * @param fft_dim dimension of 1D FFT
   * @param fft_block block size
   * @param num_fft_blocks  number of ffts
   */
  void doit(int fft_dim, int fft_block,int num_fft_blocks
      , std::vector<double>& cpu_secs, int num_tests)
  {
    int id0=num_tests*8;
    double t0=0.0,t1=0.0;
#pragma omp parallel reduction(+:t0,t1)
    {
      FFT_BENCH<real_type,real_type,false,FFTENG> single(fft_dim,fft_block,fft_dim+2,num_fft_blocks);
      single.benchmark(niters);
      t0+=single.cpu_secs(0);
      t1+=single.cpu_secs(1);
    }
    cpu_secs[id0+0]=t0;
    cpu_secs[id0+1]=t1;

    t0=0.0;t1=0.0;
#pragma omp parallel reduction(+:t0,t1)
    {
      FFT_BENCH<real_type,real_type,true,FFTENG> many(fft_dim,fft_block,fft_dim+2,num_fft_blocks);
      many.benchmark(niters);
      t0+=many.cpu_secs(0);
      t1+=many.cpu_secs(1);
    }
    cpu_secs[id0+2]=t0;
    cpu_secs[id0+3]=t1;

    t0=0.0;t1=0.0;
#pragma omp parallel reduction(+:t0,t1)
    {
      FFT_BENCH<real_type,complex_type,false,FFTENG> single(fft_dim,fft_block,fft_dim,num_fft_blocks);
      single.benchmark(niters);
      t0+=single.cpu_secs(0);
      t1+=single.cpu_secs(1);
    }
    cpu_secs[id0+4]=t0;
    cpu_secs[id0+5]=t1;

    t0=0.0;t1=0.0;
#pragma omp parallel reduction(+:t0,t1)
    {
      FFT_BENCH<real_type,complex_type,true,FFTENG> many(fft_dim,fft_block,fft_dim,num_fft_blocks);
      many.benchmark(niters);
      t0+=many.cpu_secs(0);
      t1+=many.cpu_secs(1);
    }
    cpu_secs[id0+6]=t0;
    cpu_secs[id0+7]=t1;
  }
};


int main(int argc, char** argv)
{
  mpi::environment env(argc,argv);
  mpi::communicator mycomm;
  OhmmsInfo ohmms("fft_bench",mycomm.rank(),0,1);

  ///1D FFT dimension
  int fft_dim=12288;
  ///total number of ffts
  int fft_tot=288;
  ///number of fft blocks per call
  int fft_block=1;
  int niters=10;
  unsigned plan_mode = FFTW_MEASURE;
  //using namespace fft_v0;

  bitset<MAX_FFT_ENG> test_it(0);

  //cout.setf(std::ios::scientific, std::ios::floatfield);
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-d")<a.size()) //get fft dimension
      fft_dim=atoi(argv[++ic]);
    else if(a.find("-n")<a.size()) //get how many fft to perform
      fft_tot=atoi(argv[++ic]);
    else if(a.find("-b")<a.size()) //number of simultaneous trans.
      fft_block=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("enable-fftw")<a.size())
      test_it.set(FFTW_ENG,1);
    else if(a.find("enable-essl")<a.size())
      test_it.set(FFTESSL_ENG,1);
    else if(a.find("enable-mkl")<a.size())
      test_it.set(FFTMKL_ENG,1);
    ++ic;
  }

  if(test_it.count()==0)
  {
    app_log() << "Nothing to test. Choose \n"
      << "--enable-fftw | --enable-essl | --enable-mkl"
      << endl;
    return 1;
  }

  fft_block=std::min(fft_block,fft_tot);
  fft_tot=(fft_tot/fft_block)*fft_block;

  double gflops_c=(log(static_cast<double>(fft_dim))/log(2.))*5.*fft_dim*fft_tot*1e-9;
  double gflops_r=gflops_c*0.5;

  app_log() << "#  Testing FFT dim= " << fft_dim 
    << " number of fft = " << fft_tot
    << " fft bundles = " << fft_block
    << " iterations= " << niters 
    << " MFLOPS(complex)= " << static_cast<int>(gflops_c*1000)
    << endl;

  //8=(forward+backward) x (single, multiple) x (real,complex) 
  vector<double> cpu_secs(test_it.count()*8+1,0.0);
  int num_fft_blocks=fft_tot/fft_block;

  //counter of the fftw engines
  int num_tests=0;
  vector<string> fft_names;

  Timer myclock;
#if defined(HAVE_LIBFFTW)
  if(test_it[FFTW_ENG])
  {
    test_all<double,FFTW_ENG> a("FFTW",niters);
    a.add2table(fft_names);
    a.doit(fft_dim,fft_block,num_fft_blocks,cpu_secs,num_tests);
    num_tests++;
  }
#endif
#if defined(HAVE_ESSL)
  if(test_it[FFTESSL_ENG])
  {
    test_all<double,FFTESSL_ENG> a("ESSL",niters);
    a.add2table(fft_names);
    a.doit(fft_dim,fft_block,num_fft_blocks,cpu_secs,num_tests);
    num_tests++;
  }
#endif
#if defined(HAVE_MKL)
  if(test_it[FFTMKL_ENG])
  {
    test_all<double,FFTMKL_ENG> a("MKL",niters);
    a.add2table(fft_names);
    a.doit(fft_dim,fft_block,num_fft_blocks,cpu_secs,num_tests);
    num_tests++;
  }
#endif

  cpu_secs.back()=myclock.elapsed();
  mpi::reduce(mycomm,cpu_secs);

  cout << "Number of tests " << endl;

  app_log().setf(ios::floatfield, ios::scientific);
  app_log().precision(3);
  int num_datasets=fft_names.size();
  app_log() << "Wall-clock time " << cpu_secs.back()/static_cast<double>(mycomm.size()) << endl;

  app_log() << "#fft_dim howmany block mpi openmp ";
  for(int i=0; i<fft_names.size(); ++i) app_log() << setw(12) << fft_names[i];
  app_log() << std::endl;

  double tnorm=1.0/static_cast<double>(mycomm.size()*omp_get_max_threads());
  for(int i=0; i<cpu_secs.size(); ++i) cpu_secs[i]*=tnorm;

  vector<double> gflops_fac(4);
  gflops_fac[0]=gflops_r; gflops_fac[1]=gflops_r;
  gflops_fac[2]=gflops_c; gflops_fac[3]=gflops_c;

  stringstream o;
  o.setf(ios::floatfield, ios::scientific);
  o.precision(4);
  o << fft_dim << setw(8) << fft_tot << setw(8) << fft_block << setw(4) << mycomm.size() << setw(4) << omp_get_max_threads(); 
  for(int i=0,ii=0; i<4; ++i,ii+=2)
    for(int k=0; k<num_tests; ++k)
      o << setw(12) << cpu_secs[k*8+ii] ;
  o << " f_sec " << endl;

  o << fft_dim << setw(8) << fft_tot << setw(8) << fft_block << setw(4) << mycomm.size() << setw(4) << omp_get_max_threads(); 
  for(int i=0,ii=1; i<4; ++i,ii+=2) for(int k=0; k<num_tests; ++k) o << setw(12) << cpu_secs[k*8+ii];
  o << " b_sec " << endl;


  o << fft_dim << setw(8) << fft_tot << setw(8) << fft_block << setw(4) << mycomm.size() << setw(4) << omp_get_max_threads(); 
  for(int i=0,ii=0; i<4; ++i,ii+=2)
    for(int k=0; k<num_tests; ++k) o << setw(12) << gflops_fac[i]/cpu_secs[k*8+ii];
  o << " f_gflops " << endl;

  o << fft_dim << setw(8) << fft_tot << setw(8) << fft_block << setw(4) << mycomm.size() << setw(4) << omp_get_max_threads(); 
  for(int i=0,ii=1; i<4; ++i,ii+=2)
    for(int k=0; k<num_tests; ++k) o << setw(12) << gflops_fac[i]/cpu_secs[k*8+ii];
  o << " b_gflops " << endl;

  app_log() << o.str() << endl;
  //mpi::environment::finalize();
  return 0;
}

